#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
修复LoRA表结构
"""

import sqlite3
import os

def fix_lora_table():
    """修复LoRA表结构"""
    try:
        # 数据库路径
        current_dir = os.path.dirname(os.path.abspath(__file__))
        db_path = os.path.join(current_dir, "..", "..", "flux_prompt_helper", "promptManager-server", "instance", "prompt_manager.db")
        
        if not os.path.exists(db_path):
            print(f"数据库文件不存在: {db_path}")
            return False
        
        print(f"连接数据库: {db_path}")
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()
        
        # 检查表是否存在
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='loras'")
        table_exists = cursor.fetchone()
        
        if table_exists:
            print("发现现有的loras表，需要重建...")
            
            # 备份现有数据
            cursor.execute("SELECT * FROM loras")
            existing_data = cursor.fetchall()
            print(f"备份了 {len(existing_data)} 条现有记录")
            
            # 删除旧表
            cursor.execute("DROP TABLE IF EXISTS loras")
            print("删除了旧的loras表")
        
        # 创建新的表结构
        cursor.execute("""
            CREATE TABLE loras (
                id TEXT PRIMARY KEY,
                name TEXT NOT NULL,
                path TEXT NOT NULL,
                size INTEGER DEFAULT 0,
                tags TEXT DEFAULT '[]',
                description TEXT DEFAULT '',
                images TEXT DEFAULT '[]',
                trigger TEXT DEFAULT '',
                prompt TEXT DEFAULT '',
                is_favorite BOOLEAN DEFAULT 0,
                created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
            )
        """)
        print("创建了新的loras表结构")
        
        # 如果有现有数据，尝试恢复
        if table_exists and existing_data:
            print("尝试恢复现有数据...")
            for row in existing_data:
                try:
                    # 假设旧表结构是: id, name, path, size, tags, description, images, trigger, prompt, is_favorite, created_at, updated_at
                    if len(row) >= 11:
                        cursor.execute("""
                            INSERT INTO loras 
                            (id, name, path, size, tags, description, images, trigger, prompt, is_favorite, created_at, updated_at)
                            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                        """, row)
                    print(f"恢复了记录: {row[1] if len(row) > 1 else 'Unknown'}")
                except Exception as e:
                    print(f"恢复记录失败: {e}")
        
        conn.commit()
        conn.close()
        
        print("LoRA表结构修复完成！")
        return True
        
    except Exception as e:
        print(f"修复LoRA表结构失败: {e}")
        return False

def main():
    """主函数"""
    print("开始修复LoRA表结构")
    print("=" * 40)
    
    success = fix_lora_table()
    
    if success:
        print("\n修复成功！现在可以正常保存LoRA信息了。")
    else:
        print("\n修复失败，请检查错误信息。")

if __name__ == "__main__":
    main()
